import os
import json
import threading
import random
import numpy as np
import torch
from datetime import datetime


class CommonUtils:
    """通用工具类，包含实验和超参数调优共用的方法"""
    
    @staticmethod
    def set_random_seed(seed=42):
        """
        设置全局随机种子，确保实验可复现
        
        参数:
        - seed: 随机种子值，默认42
        """
        # 设置Python内置随机数生成器
        random.seed(seed)
        # 设置NumPy随机数生成器
        np.random.seed(seed)
        # 设置PyTorch随机数生成器
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)  # 多GPU情况下
        # 确保确定性算法
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        print(f"已设置随机种子: {seed}")
    
    @staticmethod
    def get_script_dir(script_dir=None):
        """获取脚本目录"""
        return script_dir or os.path.dirname(os.path.abspath(__file__))
    
    @staticmethod
    def ensure_project_root(script_dir):
        """确保项目根目录在Python路径中"""
        import sys
        project_root = os.path.dirname(os.path.dirname(script_dir))
        if project_root not in sys.path:
            sys.path.insert(0, project_root)
        return project_root
    
    @staticmethod
    def get_run_info_file_path(script_dir):
        """获取运行信息文件路径"""
        return os.path.join(script_dir, 'run_info.json')
    
    @staticmethod
    def generate_run_name():
        """生成运行名称"""
        return f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
    
    @staticmethod
    def load_run_info(script_dir):
        """加载运行信息"""
        run_info_file = CommonUtils.get_run_info_file_path(script_dir)
        if os.path.exists(run_info_file):
            try:
                with open(run_info_file, 'r') as f:
                    return json.load(f)
            except Exception as e:
                print(f"加载运行信息失败: {e}")
        return None
    
    @staticmethod
    def save_run_info(script_dir, run_id, run_name, config):
        """保存运行信息"""
        run_info_file = CommonUtils.get_run_info_file_path(script_dir)
        run_info = {
            'run_id': run_id,
            'run_name': run_name,
            'config': config
        }
        try:
            with open(run_info_file, 'w') as f:
                json.dump(run_info, f, indent=2)
            print(f"已保存运行信息到: {run_info_file}")
            return True
        except Exception as e:
            print(f"保存运行信息失败: {e}")
            return False
    
    @staticmethod
    def cleanup_run_info(script_dir):
        """清理运行信息文件"""
        run_info_file = CommonUtils.get_run_info_file_path(script_dir)
        if os.path.exists(run_info_file):
            try:
                os.remove(run_info_file)
                print(f"已清理运行信息文件: {run_info_file}")
                return True
            except Exception as e:
                print(f"清理运行信息文件失败: {e}")
        return False


def load_config(config_path):
    """加载配置文件"""
    import yaml
    with open(config_path, 'r') as f:
        return yaml.safe_load(f)

def load_run_info(script_dir):
    """加载运行信息"""
    return CommonUtils.load_run_info(script_dir)

def save_run_info(script_dir, run_id, run_name, config):
    """保存运行信息"""
    return CommonUtils.save_run_info(script_dir, run_id, run_name, config)

def cleanup_run_info(script_dir):
    """清理运行信息文件"""
    return CommonUtils.cleanup_run_info(script_dir)


def set_random_seed(seed=42):
    """
    设置全局随机种子的便捷函数
    
    参数:
    - seed: 随机种子值，默认42
    """
    return CommonUtils.set_random_seed(seed)


def split_files_by_kfold(file_list, k=5, shuffle=True, seed=42):
    """
    将文件列表按k折交叉验证的方式划分
    
    参数:
    - file_list: 文件路径列表
    - k: 折数，默认5
    - shuffle: 是否打乱顺序，默认True
    - seed: 随机种子，默认42
    
    返回:
    - 包含k个元组的列表，每个元组为(train_files, val_files)
    """
    # 设置随机种子以保证划分结果可复现
    random.seed(seed)
    
    # 复制并打乱文件列表
    file_list_copy = file_list.copy()
    if shuffle:
        random.shuffle(file_list_copy)
    
    # 计算每折的大小
    fold_size = len(file_list_copy) // k
    folds = []
    
    # 划分数据
    for i in range(k):
        # 确定验证集的起始和结束索引
        val_start = i * fold_size
        val_end = val_start + fold_size
        if i == k - 1:  # 最后一折包含剩余所有样本
            val_end = len(file_list_copy)
        
        # 构建训练集和验证集
        val_files = file_list_copy[val_start:val_end]
        train_files = file_list_copy[:val_start] + file_list_copy[val_end:]
        
        folds.append((train_files, val_files))
    
    return folds


def async_save_file(save_func, *args, **kwargs):
    """
    在新线程中异步保存文件，避免阻塞主线程
    
    参数:
    - save_func: 用于保存文件的函数，如 torch.save
    - *args: 传递给save_func的位置参数
    - **kwargs: 传递给save_func的关键字参数
    
    返回:
    - threading.Thread: 创建的线程对象，可以用来跟踪或等待保存操作完成
    """
    
    def save_thread(save_func, args, kwargs):
        # 尝试获取文件名并提供更友好的描述
        file_name = "未知文件"
        if args and len(args) >= 2:
            # 对于torch.save，第二个参数通常是文件路径
            file_param = args[1]
            if isinstance(file_param, str):
                # 从路径中提取文件名
                base_name = os.path.basename(file_param)
 
                file_name = base_name
            elif hasattr(file_param, 'name'):
                # 如果是文件对象，获取其name属性
                file_name = file_param.name
        
        # 在开始保存前打印信息
        print(f"开始异步保存文件: {file_name}")
        
        try:
            save_func(*args, **kwargs)
            # 保存后不再打印完成信息
        except Exception as e:
            # 只在出错时打印错误信息
            print(f"异步保存文件失败: {str(e)}")
    
    # 创建并启动线程
    thread = threading.Thread(target=save_thread, args=(save_func, args, kwargs))
    thread.daemon = True  # 设置为守护线程，主线程结束时自动终止
    thread.start()
    
    return thread


def async_torch_save(obj, f, **kwargs):
    """
    异步执行torch.save操作
    
    参数:
    - obj: 要保存的对象
    - f: 文件路径或文件对象
    - **kwargs: 传递给torch.save的其他参数
    
    返回:
    - threading.Thread: 创建的线程对象
    """
    import torch
    return async_save_file(torch.save, obj, f, **kwargs)


def split_files_by_subject_kfold(file_list, k=5, data_path=None, seed=42):
    """
    按受试者对文件进行分组，并实现k折交叉验证
    特别适用于Sleep-EDF-20和Sleep-EDF-78数据集
    
    参数:
    - file_list: 文件路径列表
    - k: 折数，默认5
    - data_path: 数据路径，用于辅助判断数据集类型
    - seed: 随机种子，默认42，确保可复现性
    
    返回:
    - 包含k个元组的列表，每个元组为(train_files, val_files)
    """
    # 设置随机种子以保证划分结果可复现
    random.seed(seed)
    np.random.seed(seed)
    
    # 按受试者分组文件
    subject_files = {}
    for file_path in file_list:
        # 提取文件名
        file_name = os.path.basename(file_path)
        
        # 对于Sleep-EDF数据集，文件名格式通常为SC4xxxxx
        # 提取受试者编号（通常是文件名的前5个字符）
        if file_name.startswith('SC4'):
            subject_id = file_name[:5]  # 例如 SC401
        else:
            # 如果不是标准格式，尝试使用其他方式提取受试者ID
            # 例如取文件名中第一个下划线前的部分
            parts = file_name.split('_')
            subject_id = parts[0] if parts else file_name
        
        # 将文件添加到对应受试者的列表中
        if subject_id not in subject_files:
            subject_files[subject_id] = []
        subject_files[subject_id].append(file_path)
    
    # 将受试者的文件组转换为列表
    subject_groups = list(subject_files.values())
    
    # 对受试者组进行随机排列
    random.shuffle(subject_groups)
    
    # 计算每折的大小
    fold_size = len(subject_groups) // k
    folds = []
    
    # 划分数据
    for i in range(k):
        # 确定验证集的起始和结束索引
        val_start = i * fold_size
        val_end = val_start + fold_size
        if i == k - 1:  # 最后一折包含剩余所有样本
            val_end = len(subject_groups)
        
        # 构建训练集和验证集的受试者组
        val_groups = subject_groups[val_start:val_end]
        train_groups = subject_groups[:val_start] + subject_groups[val_end:]
        
        # 展平为文件列表
        val_files = [file for group in val_groups for file in group]
        train_files = [file for group in train_groups for file in group]
        
        folds.append((train_files, val_files))
    
    return folds